package com.ky.pro.java.mq.kmq.repository.impl;

import com.ky.pro.java.mq.kmq.core.KmqMessage;
import com.ky.pro.java.mq.kmq.repository.KmqMessageRepository;
import lombok.extern.slf4j.Slf4j;

import java.util.Collections;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;

@Slf4j
public class KmqInMemoryMessageRepository<T> implements KmqMessageRepository<T> {

    ConcurrentHashMap<String, CopyOnWriteArrayList<KmqMessage<T>>> queuePerTopic = new ConcurrentHashMap<>(64);

    /**
     * topic->{consumer->offset}
     */
    ConcurrentHashMap<String, ConcurrentHashMap<String, Integer>> offsetPerTopicPerConsumer = new ConcurrentHashMap<>(64);

    @Override
    public int saveMessageToTopic(String topic, KmqMessage<T> kmqMessage) {
        CopyOnWriteArrayList<KmqMessage<T>> queue = queuePerTopic.get(topic);
        if (queue == null) {
            throw new RuntimeException("队列" + topic + "未初始化！");
        }

        if (!queue.add(kmqMessage)) {
            throw new RuntimeException("队列" + topic + "添加元素失败！");
        }

        return queue.size();
    }

    @Override
    public int getOffset(String topic, String consumerId) {
        ConcurrentHashMap<String, Integer> offsetMap = offsetPerTopicPerConsumer.get(topic);
        if (offsetMap == null) {
            throw new RuntimeException("主题" + topic + "队列未初始化!");
        }
        Integer offset = offsetMap.get(consumerId);
        if (offset == null) {
            throw new RuntimeException("消费者" + consumerId + "未订阅主题: " + topic);
        }
        return offset.intValue();
    }

    @Override
    public boolean saveOffset(String topic, String consumerId, int offset) {
        ConcurrentHashMap<String, Integer> offsetMap = offsetPerTopicPerConsumer.get(topic);
        if (offsetMap == null) {
            log.error("主题" + topic + "队列未初始化!");
            return false;
        }
        Integer previousOffset = offsetMap.get(consumerId);
        if (previousOffset == null) {
            log.error("消费者{}未订阅主题{}！", consumerId, topic);
            return false;
        }

        log.info("消费者{}消费主题{}到位置：{}", consumerId, topic, offset);
        offsetMap.compute(consumerId, (cId, pos) -> offset);
        return true;
    }

    @Override
    public List<KmqMessage<T>> getMessage(String topic, String consumerId,
                                          int currentOffset, int batchSize) {
        CopyOnWriteArrayList<KmqMessage<T>> queue = queuePerTopic.get(topic);
        if (queue == null) {
            throw new RuntimeException("队列" + topic + "未初始化！");
        }

        int remainingMsgToConsume = queue.size() - currentOffset;

        if (remainingMsgToConsume > 0) {
            int subListToIndex = currentOffset +
                    (remainingMsgToConsume < batchSize ? remainingMsgToConsume : batchSize);

            return queue.subList(currentOffset, subListToIndex);
        }

        return Collections.emptyList();
    }

    @Override
    public boolean initQueue(String topic, int capacity) {

        if (queuePerTopic.get(topic) != null) {
            log.error("队列" + topic + "已初始化！");
            return false;
        }

        queuePerTopic.putIfAbsent(topic, new CopyOnWriteArrayList<>());
        offsetPerTopicPerConsumer.putIfAbsent(topic, new ConcurrentHashMap<>());
        return true;
    }

    @Override
    public boolean addConsumerToTopic(String topic, String consumerId) {
        ConcurrentHashMap<String, Integer> offsetMap = offsetPerTopicPerConsumer.get(topic);
        if (offsetMap == null) {
            log.error("topic:{}队列未初始化！");
            return false;
        }
        offsetMap.putIfAbsent(consumerId, new Integer(0));
        return true;
    }

    @Override
    public boolean removeConsumerFromTopic(String topic, String consumerId) {
        ConcurrentHashMap<String, Integer> offsetMap = offsetPerTopicPerConsumer.get(topic);
        if (offsetMap == null) {
            log.error("topic:{}队列未初始化！");
            return false;
        }
        Integer removedTopicOffset = offsetMap.remove(consumerId);
        log.info("消费者{}取消订阅topic: {}, 已消费offset: {}", consumerId, topic,
                removedTopicOffset);
        return true;
    }
}
